import streamlit as st
import pandas as pd
import json
from pathlib import Path
import sys
from typing import List, Dict, Any

# 添加项目路径
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))
sys.path.append(str(project_root / "src"))

st.set_page_config(
    page_title="高级搜索 - RAG-main",
    page_icon="🔍",
    layout="wide"
)

st.title("🔍 高级搜索")
st.markdown("深度搜索和分析功能")

# 初始化会话状态
if 'advanced_results' not in st.session_state:
    st.session_state.advanced_results = []
if 'comparison_results' not in st.session_state:
    st.session_state.comparison_results = {}

def load_companies():
    """加载公司列表"""
    try:
        data_path = Path("data/stock_data")
        subset_file = data_path / "subset.csv"
        
        if subset_file.exists():
            # 尝试不同的编码方式读取CSV文件
            encodings = ['utf-8', 'gbk', 'gb18030', 'latin1']
            
            for encoding in encodings:
                try:
                    df = pd.read_csv(subset_file, encoding=encoding)
                    if 'company_name' in df.columns:
                        companies = df['company_name'].unique().tolist()
                        return sorted(companies)
                    else:
                        # 如果没有company_name列，尝试查看所有列
                        st.warning(f"CSV文件中没有'company_name'列，可用列: {df.columns.tolist()}")
                        # 尝试使用第一列作为公司名
                        if len(df.columns) > 0:
                            companies = df[df.columns[0]].unique().tolist()
                            return sorted(companies)
                except UnicodeDecodeError:
                    continue
                except Exception as e:
                    st.warning(f"使用{encoding}编码读取失败: {e}")
                    continue
            
            # 如果所有编码都失败，尝试二进制方式读取并手动解析
            try:
                with open(subset_file, 'rb') as f:
                    content = f.read()
                    # 尝试检测编码
                    import chardet
                    detected = chardet.detect(content)
                    if detected['confidence'] > 0.5:
                        encoding = detected['encoding']
                        st.info(f"检测到文件编码: {encoding}，置信度: {detected['confidence']}")
                        df = pd.read_csv(subset_file, encoding=encoding)
                        if 'company_name' in df.columns:
                            companies = df['company_name'].unique().tolist()
                            return sorted(companies)
            except ImportError:
                st.warning("未安装chardet库，无法自动检测文件编码")
            except Exception as e:
                st.warning(f"尝试自动检测编码失败: {e}")
            
            st.warning("无法正确读取CSV文件，请检查文件编码")
            return []
        return []
    except Exception as e:
        st.warning(f"加载公司列表失败: {e}")
        return []

def perform_advanced_search(companies: List[str], query: str, search_params: Dict[str, Any]):
    """执行高级搜索"""
    try:
        # 导入项目模块
        from src.pipeline import Pipeline, max_config_mineru
        from src.questions_processing import QuestionsProcessor
        
        root_path = Path("data/stock_data")
        pipeline = Pipeline(root_path)
        pipeline.run_config = max_config_mineru
        
        processor = QuestionsProcessor(
            root_path=root_path,
            run_config=pipeline.run_config
        )
        
        results = {}
        for company in companies:
            with st.spinner(f"搜索 {company} 相关信息..."):
                result = processor.process_single_question(
                    company_name=company,
                    question=query
                )
                results[company] = result
        
        return results
    except Exception as e:
        st.error(f"搜索失败: {e}")
        return {}

# 侧边栏配置
with st.sidebar:
    st.header("🎯 搜索配置")
    
    # 公司选择
    companies = load_companies()
    if companies:
        selected_companies = st.multiselect(
            "选择公司（多选）",
            companies,
            default=companies[:2] if len(companies) >= 2 else companies
        )
    else:
        st.warning("未找到公司数据")
        selected_companies = []
    
    # 搜索模式
    search_mode = st.selectbox(
        "搜索模式",
        ["单一查询", "批量查询", "对比分析"],
        help="选择搜索模式"
    )
    
    # 高级参数
    st.subheader("高级参数")
    top_k = st.slider("检索文档数量", 3, 20, 6)
    similarity_threshold = st.slider("相似度阈值", 0.0, 1.0, 0.5, 0.1)
    use_reranking = st.checkbox("启用 LLM 重排序", True)

# 主内容区域
tab1, tab2 = st.tabs(["🔍 多公司搜索", "📊 对比分析"])

with tab1:
    st.header("多公司搜索")
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # 查询输入
        if search_mode == "单一查询":
            query = st.text_area(
                "输入查询问题",
                height=100,
                placeholder="例如：公司的主要业务是什么？"
            )
            
            if st.button("🔍 开始搜索", type="primary") and query and selected_companies:
                search_params = {
                    "top_k": top_k,
                    "similarity_threshold": similarity_threshold,
                    "use_reranking": use_reranking
                }
                
                results = perform_advanced_search(selected_companies, query, search_params)
                st.session_state.advanced_results = results
        
        elif search_mode == "批量查询":
            st.subheader("批量查询")
            
            # 预设查询模板
            query_templates = [
                "公司的主要业务是什么？",
                "公司的营收情况如何？",
                "公司面临的主要风险有哪些？",
                "公司的研发投入情况？",
                "公司的市场地位如何？"
            ]
            
            selected_queries = st.multiselect(
                "选择查询问题",
                query_templates,
                default=query_templates[:3]
            )
            
            if st.button("🚀 批量搜索", type="primary") and selected_queries and selected_companies:
                batch_results = {}
                progress_bar = st.progress(0)
                
                total_queries = len(selected_queries)
                
                for i, query in enumerate(selected_queries):
                    search_params = {
                        "top_k": top_k,
                        "similarity_threshold": similarity_threshold,
                        "use_reranking": use_reranking
                    }
                    
                    results = perform_advanced_search(selected_companies, query, search_params)
                    batch_results[query] = results
                    
                    progress_bar.progress((i + 1) / total_queries)
                
                st.session_state.advanced_results = batch_results
                st.success("批量搜索完成！")
    
    with col2:
        st.subheader("搜索统计")
        
        if selected_companies:
            st.metric("选中公司数", len(selected_companies))
        
        if st.session_state.advanced_results:
            if isinstance(st.session_state.advanced_results, dict):
                total_results = len(st.session_state.advanced_results)
                st.metric("搜索结果数", total_results)
    
    # 显示搜索结果
    if st.session_state.advanced_results:
        st.subheader("搜索结果")
        
        if search_mode == "单一查询":
            # 单一查询结果显示
            for company, result in st.session_state.advanced_results.items():
                with st.expander(f"📊 {company}", expanded=True):
                    
                    # 显示答案
                    if 'answer' in result:
                        st.markdown("**生成答案:**")
                        st.info(result['answer'])
                    
                    # 显示检索文档
                    if 'retrieved_documents' in result:
                        st.markdown("**相关文档:**")
                        docs = result['retrieved_documents']
                        
                        for i, doc in enumerate(docs[:3]):  # 显示前3个
                            st.markdown(f"**文档 {i+1}** (相似度: {doc.get('score', 0):.3f})")
                            content = doc.get('text', '')[:300] + "..."
                            st.write(content)

with tab2:
    st.header("📊 对比分析")
    
    if len(selected_companies) >= 2:
        # 对比查询
        comparison_query = st.text_input(
            "对比查询问题",
            placeholder="例如：比较各公司的盈利能力"
        )
        
        if st.button("🔄 开始对比", type="primary") and comparison_query:
            search_params = {
                "top_k": top_k,
                "similarity_threshold": similarity_threshold,
                "use_reranking": use_reranking
            }
            
            comparison_results = perform_advanced_search(selected_companies, comparison_query, search_params)
            st.session_state.comparison_results = comparison_results
        
        # 显示对比结果
        if st.session_state.comparison_results:
            st.subheader("对比结果")
            
            # 创建对比表格
            comparison_data = []
            for company, result in st.session_state.comparison_results.items():
                comparison_data.append({
                    "公司": company,
                    "答案": result.get('answer', '')[:100] + "...",
                    "相关文档数": len(result.get('retrieved_documents', [])),
                    "平均相似度": sum(doc.get('score', 0) for doc in result.get('retrieved_documents', [])) / max(len(result.get('retrieved_documents', [])), 1)
                })
            
            df_comparison = pd.DataFrame(comparison_data)
            st.dataframe(df_comparison, use_container_width=True)
    else:
        st.warning("对比分析需要选择至少2家公司")